#作者：神奇的小阿飞
import bpy
import bmesh
import numpy as np
from mathutils import Vector;
import math;

maxValue = 5;
minValue = -5;
frameNumber = 5;
frameStart = 0;

ExportNormal = True;
ExportVertex = True;

rangeValue = maxValue-minValue;

# Get the active mesh
me = bpy.context.object.data;
object = bpy.context.object;
dg = bpy.context.evaluated_depsgraph_get();

def to01(v):
    if(v>maxValue):
        return 1;
    if(v<minValue):
        return 0;
    return (v-minValue)/rangeValue;


def createImage(name,w,h,alpha):
    img = bpy.data.images.get(name);
    if(img!=None):
        bpy.data.images.remove(img);
        img = None;
    img = bpy.data.images.new(name, width=w, height=h, alpha=alpha, float_buffer=True);
    #img.colorspace_settings.name='Linear'
    return img;

if(dg==None):
    print("dg is none!!");

bpy.context.scene.frame_set(0);
bm = bmesh.new();
bm.from_object(object,dg);
bmesh.ops.triangulate(bm, faces=bm.faces);


vertNum =len(bm.faces)*3

unitOffset = 1.0 /vertNum;
half_offset = unitOffset*(0.5);

#set uv
id = 0;
#bpy.ops.mesh.uv_texture_add();
uv_layer = bm.loops.layers.uv.active;

for f in bm.faces:
    for loop_data in f.loops:
        #uv_data = loop_data[0][uv_layer].uv;
        loop_data[uv_layer].uv.x = id*unitOffset+half_offset;
        loop_data[uv_layer].uv.y = 0.0;
        id = id +1;
bm.to_mesh(me);       

width = vertNum;
height = frameNumber;

print("Verts Number = %s" % vertNum);
print("width = %s" % width);
print("height = %s" % height);

if(ExportVertex):
    vertex_data = np.ones(width*height*4);
    vimg = createImage("Vertex_Texture",width,height,False);
if(ExportNormal):
    normal_data = np.ones(width*height*4);
    nimg = createImage("Normal_Texture",width,height,False);

index = 0;
 
   
# Modify the BMesh, can do anything here...
for frame in range(height):
    for f in bm.faces:
        for v in f.verts:         
            #save vertex data
            if(ExportVertex):
                wm = object.matrix_world.copy()
                pos = object.matrix_world@v.co;
                #print(pos);
                #vector -> [0,1]
                x = to01(pos.x);
                y = to01(pos.y);
                z = to01(pos.z);
                vertex_data[index  ] = x;
                vertex_data[index+1] = y;
                vertex_data[index+2] = z;
                vertex_data[index+3] = 1;
            #save normal data
            if(ExportNormal):
                vn = v.normal;
                normal_data[index  ] = (vn.x+1)/2.0;
                normal_data[index+1] = (vn.y+1)/2.0;
                normal_data[index+2] = (vn.z+1)/2.0;
                normal_data[index+3] = 1;
                        
            index =index+4;
        
    bm = bmesh.new();
    bm.from_object(object,dg);
    bmesh.ops.triangulate(bm, faces=bm.faces);
    #Jump Frame
    bpy.context.scene.frame_set(frameStart + frame);

#save data    
if(ExportVertex):
    vimg.pixels[:] = vertex_data;
if(ExportNormal):    
    nimg.pixels[:] = normal_data;
    
#img.save_as(filename="D:/img.png");
del bm
